package com.milvus.featureExtractor;

import com.milvus.imagesearch.MilvusManager;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.nd4j.linalg.api.ndarray.INDArray;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;

/**
 * 视频帧提取特征向量
 */
public class VideoFrameFeatureExtractor {



    public  boolean extractFeatures( String collectionName, String videoPath) throws IOException {
        MilvusManager milvusManager = new MilvusManager();
        ImageFeatureExtractor extractor = new ImageFeatureExtractor();

        String outputDir = createOutPutDirectory(videoPath); // 输出目录
        int frameRate = 1; // 每秒抽取的帧数
        try {


            FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(videoPath);
            grabber.start();

            int frameNumber = 0;
            int frameCount = 0;
            Frame frame;
            Java2DFrameConverter converter = new Java2DFrameConverter();

            while ((frame = grabber.grab()) != null) {
                if (frame.image != null) { // 检查帧是否有效
                    BufferedImage bufferedImage = converter.convert(frame);
                    if (bufferedImage != null) { // 检查转换后的图像是否有效
                        File outputFile = new File(outputDir + "frame_" + frameCount + ".png");
                        ImageIO.write(bufferedImage, "png", outputFile);

                        INDArray features = extractor.extractFeatures(outputFile);
                        milvusManager.insertData( collectionName,features,outputFile.getName());
                        frameCount++;
                    } else {
                        System.out.println("Failed to convert frame " + frameNumber);
                    }
                } else {
                    System.out.println("Invalid frame " + frameNumber);
                }
                frameNumber++;
            }
            grabber.stop();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return true;
    }

    private static String createOutPutDirectory(String videoPath) {
        File file1 = new File(videoPath);
        String outputDir = file1.getParent() + "\\" + file1.getName().split("\\.")[0]+"\\";
        //检测文件夹是否存在,不存在创建文件夹,存在清空文件夹
        File file = new File(outputDir);
        if (!file.exists()) {
            file.mkdirs();
        } else {
            for (File f : file.listFiles()) {
                f.delete();
            }
        }
        return  outputDir;
    }
}
